import datetime
import numpy as np
import matplotlib.colors as colors
import matplotlib.finance as finance
import matplotlib.dates as mdates
import matplotlib.ticker as mticker
import matplotlib.mlab as mlab
import matplotlib.pyplot as plt
import matplotlib.font_manager as font_manager
import pyodbc, datetime
from numpy import array

#startdate = datetime.date(2010,8,10)

def get_historical_database(ticker, startdate, enddate):
    cnxn = pyodbc.connect('DSN=Stock')
    cursor = cnxn.cursor()
    query = "select [Time], [Open], [High], [Low], [Close], [Volume], [WAP]\
            from Stock.dbo.stock_fivemin\
            where Symbol = '" + ticker + "' and [Time] between '" + str(startdate) + "' and '" + str(enddate) + "'\
            order by [Time]"
    cursor.execute(query)
    result = cursor.fetchall()
    #_close = []
    #for r in result:
        #print r
    cnxn.close()
    return result

def moving_average(x, n, type='simple'):
    """
    compute an n period moving average.
    type is 'simple' | 'exponential'
    """
    x = np.asarray(x)
    if type=='simple':
        weights = np.ones(n)
    else:
        weights = np.exp(np.linspace(-1., 0., n))

    weights /= weights.sum()

    a =  np.convolve(x, weights, mode='full')[:len(x)]
    a[:n] = a[n]
    return a


def moving_average2(x, n, type='simple'):
    """
    computer a n period moving avergage.
    type is 'simple' | 'exponential'
    should raise error when len(x) != n
    """
    if type=='simple':
        weights = np.ones(n)
    else:
        weights = np.exp(np.linspace(-1., 0., n))

    weights /= weights.sum()

    return np.dot(x, weights)


def relative_strength(prices, n=14):
    """
    compute the n period relative strength indicator
    """

    deltas = np.diff(prices)
    seed = deltas[:n+1]
    up = seed[seed>=0].sum()/n
    down = -seed[seed<0].sum()/n
    rs = up/down
    rsi = np.zeros_like(prices)
    rsi[:n] = 100. - 100./(1.+rs)

    for i in range(n, len(prices)):
        delta = deltas[i-1] # cause the diff is 1 shorter

        if delta>0:
            upval = delta
            downval = 0.
        else:
            upval = 0.
            downval = -delta

        up = (up*(n-1) + upval)/n
        down = (down*(n-1) + downval)/n

        rs = up/down
        rsi[i] = 100. - 100./(1.+rs)

    return rsi

def relative_strength2(prices, n=14):
    """
    compute the n period relative strength indicator
    """

    deltas = np.diff(prices)
    seed = deltas[:n+1]
    up = seed[seed>=0].sum()/n
    down = -seed[seed<0].sum()/n
    if down == 0.0:
        return 100.
    else:
        rs = up/down
        return 100. - 100./(1.+rs)

def moving_average_convergence(x, nslow=26, nfast=12):
    """
    compute the MACD (Moving Average Convergence/Divergence) using a fast and slow exponential moving avg'
    return value is emaslow, emafast, macd which are len(x) arrays
    """
    emaslow = moving_average(x, nslow, type='exponential')
    emafast = moving_average(x, nfast, type='exponential')
    return emaslow, emafast, emafast - emaslow

def williams_percent_r(_close, _high, _low, p=14):
    williams_r = np.zeros_like(_close)

    for i in range(p):
        hh = _high[0:i+1].max()
        ll = _low[0:i+1].min()
        williams_r[i] = -100.*(hh - _close[i])/(hh-ll)

    for i in range(p, len(_close)):
        hh = _high[i-p:i+1].max()
        ll = _low[i-p:i+1].min()
        williams_r[i] = -100.*(hh - _close[i])/(hh-ll)

    return williams_r

def williams_percent_r2(_close, _high, _low, p=14):
    return -100.*(_high.max() - _close[len(_close)-1])/(_high.max() - _low.min())

def true_strength_indicator2(prices, nslow=14, nfast=7):
    """
    length of prices should be at least nslow+nfast
    delta should be at least nslow+nfast-1
    """
    delta = np.diff(prices) #length reduce by one
    ema1 = moving_average(delta, nslow, 'exponential')
    ema2 = moving_average(np.abs(delta), nslow, 'exponential')
    fast_ema2 = moving_average2(ema2[-nfast:], nfast, 'exponential')
    if fast_ema2 == 0.0:
        return 100
    else:
        fast_ema1 = moving_average2(ema1[-nfast:], nfast, 'exponential')
        return 100.*fast_ema1/fast_ema2

def ergodic_candlestick_oscillator2(_close, _open, _high, _low, r=14, s=7):
    """
    length of prices should be at least r+s
    delta should be at least r+s-1
    if fast_emahl = 0.0, it's hard to define, it's 100 or 0
    """
    emaco = moving_average(_close-_open, r, 'exponential')
    emahl = moving_average(_high-_low, r, 'exponential')
    fast_emahl = moving_average2(emahl[-s:], s, 'exponential')
    if fast_emahl == 0.0:
        return 100
    else:
        fast_emaco = moving_average2(emaco[-s:], s, 'exponential')
        return 100.*fast_emaco/fast_emahl

    #return 100.*moving_average2(moving_average(delta, nslow, 'exponential'), nfast, 'exponential')/moving_average2(moving_average(np.abs(delta), nslow, 'exponential'), nfast, 'exponential')

def update_fifo_queue(xlist, nvalue, length=50):
    if len(xlist) < length:
        xlist.append(nvalue)
    else:
        xlist.pop(0)

if __name__ == '__main__':
    initialCapital = 10000.
    openrangePeriod = 15
    stoploss, stopprofit = -1.0, 1.0
    #cutloss, takeprofit = -0.01, 0.003
    transcost = 0.000
    commission = 2.01

    longOnly = False
    beginTime = '09:30:00'
    endTime = '15:30:00'
    ticks = 2
    barExit = 1
    #goLong = False

    totalPnl = 0.

    #date = ''
    #15or_h, 15or_l = 0., 0.
    cnxn = pyodbc.connect('DSN=Stock')
    cursor = cnxn.cursor()
    query = "select convert(varchar(10), [Time], 111), [High], [Low]\
            from Stock.dbo.esh1_fifteenmin\
            where convert(varchar(10), [Time], 8) = '09:30:00'\
            and convert(varchar(10), [Time], 111) < '2011/03/18'\
            order by [Time] asc"
    cursor.execute(query)
    result = cursor.fetchall()
    _date, _high, _low = [r[0] for r in result], [r[1] for r in result], [r[2] for r in result]
    #print _date
    #print _high
    #print _low

    for i in range(len(_date)):
        date = _date[i]
        or15_h = float(_high[i])
        or15_l = float(_low[i])

        query = "select [Time], [Open], [High], [Low], [Close]\
                from Stock.dbo.esh1_onemin_clean\
                where convert(varchar(10), [Time], 111) = '" + date + "' \
                and convert(varchar(10), [Time], 8) between '" + beginTime + "' and '" + endTime + "' \
                order by [Time] asc"
        #query = "select [Time], [Open], [High], [Low], [Close]\
                #from Stock.dbo.esh1_onemin_clean\
                #where convert(varchar(10), [Time], 111) = '" + date + "' \
                #order by cast(floor(cast([Time] as float)) as datetime)"
        cursor.execute(query)
        result = cursor.fetchall()
        _time_t = [r[0] for r in result]
        _open_t = [float(r[1]) for r in result]
        _high_t = [float(r[2]) for r in result]
        _low_t = [float(r[3]) for r in result]
        _close_t = [float(r[4]) for r in result]

        # try 15 minute open range first
        size_t = 1
        breakOut = False
        entryP, exitP = 0., 0.
        lossStopped, profitStopped = False, False
        pnl_t = 0.
        for j in range(openrangePeriod,150):
            if _high_t[j] >= or15_h+.25*ticks and not breakOut:
                breakOut = True
                entryP = -max(_low_t[j], or15_h+.25*ticks)
                entryBar = j
                #print "%s long with price %.2f" % (_time_t[j], entryP)
            if _low_t[j] <= or15_l-.25*ticks and not breakOut:
                breakOut = True
                entryP = min(_high_t[j], or15_l-.25*ticks)
                #print "%s short with price %.2f" % (_time_t[j], entryP)
                entryBar = j

            # long position
            if breakOut and entryP<0 and j>entryBar and j<=entryBar+barExit:
                # if stop loss and stop profit conditions met in the same bar, stop loss disregarding which one comes earlier
                if _low_t[j]+entryP <= stoploss:
                    pnl = 50*stoploss - 2*size_t*commission
                    totalPnl += pnl
                    lossStopped = True
                    print "%s long with stop loss" % _time_t[j]
                    print "accumulative pnl is %.2f" % totalPnl
                    break
                if _low_t[j]+entryP >= stopprofit:
                    pnl = 50*stoploss - 2*size_t*commission
                    totalPnl += pnl
                    profitStopped = True
                    print "%s long with stop profit" % _time_t[j]
                    print "accumulative pnl is %.2f" % totalPnl
                    break
                if not profitStopped and not lossStopped and j==entryBar+barExit:
                    pnl = 50*(_close_t[j]+entryP) - 2*size_t*commission
                    totalPnl += pnl
                    print "%s long with pnl = %.2f" % (_time_t[j], pnl)
                    print "accumulative pnl is %.2f" % totalPnl
                    break


            # short position
            if breakOut and entryP>0 and j>entryBar and j<=entryBar+barExit:
                # if stop loss and stop profit conditions met in the same bar, stop loss disregarding which one comes earlier
                if -_high_t[j]+entryP <= stoploss:
                    pnl = 50*stoploss - 2*size_t*commission
                    totalPnl += pnl
                    lossStopped = True
                    print "%s short with stop loss" % _time_t[j]
                    print "accumulative pnl is %.2f" % totalPnl
                    break
                if -_high_t[j]+entryP >= stopprofit:
                    pnl = 50*stoploss - 2*size_t*commission
                    totalPnl += pnl
                    profitStopped = True
                    print "%s short with stop profit" % _time_t[j]
                    print "accumulative pnl is %.2f" % totalPnl
                    break
                if not profitStopped and not lossStopped and j==entryBar+barExit:
                    pnl = 50*(-_close_t[j]+entryP) - 2*size_t*commission
                    totalPnl += pnl
                    print "%s short with pnl = %.2f" % (_time_t[j], pnl)
                    print "accumulative pnl is %.2f" % totalPnl
                    break

    cnxn.close()

    print "total pnl is %.2f" % totalPnl
    """
    barData = get_historical_database(ticker, startdate, enddate)

    #for r in result:
        #_datetime.append(r[0])
        #_open.append(r[1])
        #_high.append(r[2])
        #_low.append(r[3])
        #_close.append(r[4])
        #_volume.append(r[5])

    #indicators = np.array()
    #weights = np.arrap()

    listlen = 50
    investment = 100000.
    buypower = 100000.
    position = 0
    avgPrice = 0.

    for r in barData:

        if r[0].hour == 9 and r[0].minute == 30:
            _close, _datetime, _low, _high, _open, _volume, _wap = [], [], [], [], [], [], []

        update_fifo_queue(_datetime, r[0], listlen)
        update_fifo_queue(_open, r[1], listlen)
        update_fifo_queue(_high, r[2], listlen)
        update_fifo_queue(_low, r[3], listlen)
        update_fifo_queue(_close, r[4], listlen)
        update_fifo_queue(_volume, r[5], listlen)
        update_fifo_queue(_wap, r[6], listlen)

        if len(_datetime) == listlen:
            dateTime, closeP, lowP, highP, openP, volume = array(_datetime), array(_close), array(_low), array(_high), array(_open), array(_volume)
            #dateTime, closeP, lowP, highP, openP, volume = zeros(listlen), zeros(listlen), zeros(listlen), zeros(listlen), zeros(listlen), zeros(listlen)

            rsi1_p, rsi2_p, rsi3_p, rsi4_p, rsi5_p = 2, 7, 9, 14, 21
            #rsi = relative_strength(closeP[listlen-rsi_p:listlen+1], rsi_p)
            rsi1 = relative_strength2(closeP[-rsi1_p:], rsi1_p)
            rsi2 = relative_strength2(closeP[-rsi2_p:], rsi2_p)
            rsi3 = relative_strength2(closeP[-rsi3_p:], rsi3_p)
            rsi4 = relative_strength2(closeP[-rsi4_p:], rsi4_p)
            rsi5 = relative_strength2(closeP[-rsi5_p:], rsi5_p)

            mr1_p, mr2_p, mr3_p = 7, 14, 21
            #mr = williams_percent_r(closeP[listlen-mr_p:listlen+1], highP[listlen-mr_p:listlen+1], lowP[listlen-mr_p:listlen+1], mr_p)
            mr1 = 100.+williams_percent_r2(closeP[-mr1_p:], highP[-mr1_p:], lowP[-mr1_p:], mr1_p)
            mr2 = 100.+williams_percent_r2(closeP[-mr2_p:], highP[-mr2_p:], lowP[-mr2_p:], mr2_p)
            mr3 = 100.+williams_percent_r2(closeP[-mr3_p:], highP[-mr3_p:], lowP[-mr3_p:], mr3_p)

            tsi1_p1, tsi1_p2 = 4,2
            tsi2_p1, tsi2_p2 = 8,2
            tsi3_p1, tsi3_p2 = 8,4
            tsi4_p1, tsi4_p2 = 14,7
            tsi5_p1, tsi5_p2 = 21,14
            tsi1 = (100.+true_strength_indicator2(closeP[-tsi1_p1-tsi1_p2:], tsi1_p1, tsi1_p2))/2
            tsi2 = (100.+true_strength_indicator2(closeP[-tsi2_p1-tsi2_p2:], tsi2_p1, tsi2_p2))/2
            tsi3 = (100.+true_strength_indicator2(closeP[-tsi3_p1-tsi3_p2:], tsi3_p1, tsi3_p2))/2
            tsi4 = (100.+true_strength_indicator2(closeP[-tsi4_p1-tsi4_p2:], tsi4_p1, tsi4_p2))/2
            tsi5 = (100.+true_strength_indicator2(closeP[-tsi5_p1-tsi5_p2:], tsi5_p1, tsi5_p2))/2

            eco1_p1, eco1_p2 = 4,2
            eco2_p1, eco2_p2 = 6,3
            eco3_p1, eco3_p2 = 8,4
            eco4_p1, eco4_p2 = 11,4
            eco5_p1, eco5_p2 = 14,5
            eco1 = (100.+ergodic_candlestick_oscillator2(closeP[-eco1_p1-eco1_p2:], openP[-eco1_p1-eco1_p2:], highP[-eco1_p1-eco1_p2:], lowP[-eco1_p1-eco1_p2:], eco1_p1, eco1_p2))/2
            eco2 = (100.+ergodic_candlestick_oscillator2(closeP[-eco2_p1-eco2_p2:], openP[-eco2_p1-eco2_p2:], highP[-eco2_p1-eco2_p2:], lowP[-eco2_p1-eco2_p2:], eco2_p1, eco2_p2))/2
            eco3 = (100.+ergodic_candlestick_oscillator2(closeP[-eco3_p1-eco3_p2:], openP[-eco3_p1-eco3_p2:], highP[-eco3_p1-eco3_p2:], lowP[-eco3_p1-eco3_p2:], eco3_p1, eco3_p2))/2
            eco4 = (100.+ergodic_candlestick_oscillator2(closeP[-eco4_p1-eco4_p2:], openP[-eco4_p1-eco4_p2:], highP[-eco4_p1-eco4_p2:], lowP[-eco4_p1-eco4_p2:], eco4_p1, eco4_p2))/2
            eco5 = (100.+ergodic_candlestick_oscillator2(closeP[-eco5_p1-eco5_p2:], openP[-eco5_p1-eco5_p2:], highP[-eco5_p1-eco5_p2:], lowP[-eco5_p1-eco5_p2:], eco5_p1, eco5_p2))/2

            w_rsi1, w_rsi2, w_rsi3, w_rsi4, w_rsi5 = 1., 1., 1., 1., 1.
            #indicator = w_rsi1*rsi1 + w_rsi2*rsi2 + w_rsi3*rsi3 + w_rsi4*rsi4 + w_rsi5*rsi5
            #print indicator
            w_tsi1, w_tsi2, w_tsi3, w_tsi4, w_tsi5 = 1., 1., 1., 1., 1.

            indicators = array([rsi1, rsi2, rsi3, rsi4, rsi5, mr1, mr2, mr3, tsi1, tsi2, tsi3, tsi4, tsi5, eco1, eco2, eco3, eco4, eco5])
            weights = np.ones_like(indicators)/indicators.size

            #print '%.2f %.2f %.2f' % (mr1, mr2, mr3)
            my_indicator = np.dot(indicators, weights)

            #print r[0], r[1], r[4], r[6], my_indicator



            if my_indicator < oversold and position == 0 and (r[0].hour < 15 or (r[0].hour == 15 and r[0].minute <= 30)):
                goLong = True
                continue

            if position > 0:
                # use open or use wap or use last minute last or last minute wap
                PNL = (r[6] - avgPrice)/avgPrice - trans_cost
                #PNL = (r[1] - avgPrice)/avgPrice - trans_cost
                #print unrealized_loss, cutloss, takeprofit
                #if my_indicator > overbought or (r[0].hour == 15 and r[0].minute == 55):
                if my_indicator > overbought or PNL < cutloss or PNL > takeprofit or (r[0].hour == 15 and r[0].minute == 55):
                    buypower = position*r[1]
                    print 'liquid %d shares at %s with avg price %.2f, PNL is %.3f%%' % (position, r[0], r[6], PNL*100.)
                    if PNL > 0:
                        win.append(PNL*100.)
                    else:
                        loss.append(PNL*100.)
                    position = 0

            if goLong:
                position = int(buypower/r[6])
                avgPrice = r[6]
                #buypower = 0.
                goLong = False
                print 'bought %d shares at %s with avg price %.2f' % (position, r[0], r[6])
            #print '%.2f,%.2f,%.2f' % (my_indicator, r[1], r[4])
            #indicators =

            #indicator = w_rsi1*rsi1 + w_rsi2*rsi2 + w_rsi3*rsi3 + w_rsi4*rsi4 + w_rsi5*rsi5
            #print "%.2f %.2f %.2f %.2f %.2f %.2f %.2f %.2f %.2f %.2f %.2f" % (rsi1, rsi2, rsi3, rsi4, rsi5, mr2, tsi1, tsi2, tsi3, tsi4, tsi5)

            #print "%.2f %.2f %.2f %.2f %.2f %.2f %.2f %.2f %.2f %.2f %.2f" % (rsi1, rsi2, rsi3, rsi4, rsi5, mr2, tsi1, tsi2, tsi3, tsi4, tsi5)
            #print tsi1

    print 'maximum drop down is %.2f%%' % (min(loss))
    print 'maximum gain is %.2f%%' % (max(win))
    print 'total number of trades is %d' % (len(loss)+len(win))
    print 'winning rate is %.2f%%' % (100.*len(win)/(len(loss)+len(win)))
    print 'average profit is %.2f%%' % ((sum(loss)+sum(win))/(len(loss)+len(win)))
    """
